from sentence_transformers import SentenceTransformer
import chromadb
from llm.doubao import doubao_qa
from llm.local import ollama_qa
import json

model = SentenceTransformer("all-MiniLM-L6-v2")

client = chromadb.PersistentClient(path="./chromadb_data")
collection = client.get_or_create_collection("rag")


def get_query_embedding(query):
    return model.encode(query).tolist()


def abstractify_query(original_query):
    abstraction_prompt = f"""
    请对以下的用户查询进行抽象化的转换，将其转化为更高层次的抽象问题
    技术原理：
    1. 细节剥离：去除查询中的具体的细节信息(比如说时间、地点、具体的人名等)
    2. 概念抽象：将具体的问题转化为抽象的概念
    3. 范围扩展： 扩大查询 的语义范围

    用户查询:{original_query}

    请生成一个抽象化的查询版本：要求
    1. 保持原始查询的核心意图
    2. 去除过于具体的细节信息
    3. 使用更通用的概念和表达
    4. 扩展大查询的语义范围
    5. 确保抽象化的查询仍然有意义

    要求以JSON格式返回结果：
    {
        {
            "abstract_query":"抽象化后的表达",
            "abstraction_reason":"抽象化的原因说明"
        }
    }
    """
    abstract_result = ollama_qa(abstraction_prompt)
    jsons_tart = abstract_result.find("{")
    jsons_end = abstract_result.rfind("}") + 1
    json_str = abstract_result[jsons_tart:jsons_end]
    parsed_result = json.loads(json_str)
    abstract_query = parsed_result.get("abstract_query", "")
    abstraction_reason = parsed_result.get("abstraction_reason", "")
    return abstract_query, abstraction_reason


def retrieve_related_chunks_by_query(query, n_results=3):
    query_embedding = get_query_embedding(query)
    results = collection.query(query_embeddings=[query_embedding], n_results=n_results)
    related_chunks = results.get("documents")
    if not related_chunks or not related_chunks[0]:
        exit(1)
    return related_chunks[0]


def hybird_retrieve_and_answer(query, abstract_query, abstraction_reason):
    # 1.使用抽象化的查询进行广泛检索
    abstract_chunks = retrieve_related_chunks_by_query(abstract_query, n_results=5)
    # 2.使用原始查询精确的检索
    original_chunks = retrieve_related_chunks_by_query(query, n_results=3)
    # 3.合并检索结果并去重
    unique_chunks = list(set(abstract_chunks + original_chunks))
    # 4.构建提示词
    context = "\n".join(unique_chunks)
    prompt = f"""
    根据以下的检索到相关信息，回答用户的原始问题

    检索到的相关信息:
    {context}

    用户的原始问题:{query}

    注意：系统使用了抽象化查询转换技术来扩大检索范围。
    - 原始查询：{query}
    - 抽象化查询：{abstract_query}
    - 抽象化原因：{abstraction_reason}
    
    请提供一个准确、全面的回答，确保：
    1. 直接回答用户的原始问题
    2. 充分利用检索到的相关信息
    3. 如果信息不足，请明确指出
    4. 保持回答的准确性和相关性

    """
    return ollama_qa(prompt)


if __name__ == "__main__":
    query = "如何做红烧肉?"
    # 抽象化查询转换
    abstract_query, abstraction_reason = abstractify_query(query)
    print("abstract_query", abstract_query)
    print("abstraction_reason", abstraction_reason)
    # 混合检索和答案生成
    final_answer = hybird_retrieve_and_answer(query, abstract_query, abstraction_reason)
    print("final_answer", final_answer)
